#!/usr/bin/env Rscript
args = commandArgs(trailingOnly=TRUE)

library(Rcpp)
require(Matrix)
require(ranger)
require(grf)

s <- as.integer(args[1])


# Error draws
E <- 50

# Number of trees
ntrees <- 250

# HEADER ####
# setwd(args[1]) # a

base_path <- "model_output"

Rcpp::sourceCpp('code/dynamic model/functions.cpp')
source("code/dynamic model/R_functions.R")


# 2.1 Reload data ####
full_data <- readRDS("data/tmp/model_input_data.RDS")
full_data <- full_data[order(full_data$hid, full_data$year),]


rrf <- readRDS(sprintf("%s/random_forest_weights.RDS", base_path))
state.vars <- names(rrf$X.orig)


sel <- full_data$stay_length == 0 & full_data$tenure == 1
sizes_sqm <- round(aggregate(
  full_data$rent0[sel] / full_data$rent_psqm_adj[sel],
  by=list(size = full_data$size[sel]), FUN = mean
)[,2])

data <- full_data[ ! is.na(full_data$depvar_adj),]
data$choice <- data$depvar_adj
data <- data[order(data$hid, data$year),]
hid_years <- paste(data$hid, data$year, sep=":")


J <- max(data$depvar_adj)

X <- data[,state.vars]

Y = as.integer(data$choice ==  41)
ranger.data <- cbind(
  Y = Y, X
)

#
# Other parameters
beta <- .95
tol.theta <- 1e-5
tol.pi <- 1e-5

g = 0 # rent growth rate

uvars <- c(
  uvars_c <- c(
    'log_dispinc' # = log(income - rent - 2.5 * size) #
    ,'hqual'
    ,'log_size_ratio'
    ,'stay_length'
  ),
  uvars_s <- c(
    'move'
    ,"move_age"
    ,"move_age_sq"
    ,"move_two_adults"
    ,"move_kids"
  )
)
#
tvars <- c(
  "intercept",
  "log_income",
  "gwealth_imp",
  "gwealth_imp_sq",
  "age",
  "age_sq",
  "two_adults",
  "kids"
)

#
nvars_u_c = length(uvars_c) # number of constant variables in the flow utility function
nvars_u_s = length(uvars_s) # number of type-specific variables in the flow utility function
nvars_v = length(tvars)

go_file <- sprintf("%s/iter/go_files/go_file%d", base_path, s)
keep_going_file <- sprintf("%s/iter/go_files/keep_going", base_path)
counter <- 0
seed <- 381271
while(file.exists(keep_going_file)) {
  if(file.exists(go_file)) {

    if(counter == 1) seed <- 381271 #
    counter <- counter + 1

    # create ranger.rrf
    q_s <- readRDS(file=sprintf("%s/iter/q%d.RDS", base_path, s))
    q_s <- q_s / sum(q_s) * length(q_s)
    rrf <- regression_forest(
      X=X,
      Y=Y,
      num.trees = ntrees,
      mtry=rrf$tunable.params$mtry,
      alpha=rrf$tunable.params$alpha,
      sample.fraction=rrf$tunable.params$sample.fraction,
      imbalance.penalty = rrf$tunable.params$imbalance.penalty,
      sample.weights = q_s
    )

    v <- pre_compute_v_s(
      x = data, # data for one household, all periods (excluding the last period, where the choice is not yet observed)
      state.vars = state.vars, # names of state variables
      beta = beta, # discount factor
      nvars_u_c = nvars_u_c, # number of constant variables in the flow utility function
      nvars_u_s = nvars_u_s, # number of type-specific variables in the flow utility function
      nvars_v = nvars_v, # number of variables in the terminal utility functions
      compute_psi = T, # set to F if just changing the utility function (skips psi function evaluation)
      sizes_sqm = sizes_sqm,
      S = 1, # number of unobserved states
      E = E,
      save_W = F,#save_W,
      base_path = base_path,
      version = version,
      seed = seed,
      g = g
    )
    file.remove(go_file)
    saveRDS(v[,c(1,2,4,ncol(v)-1)], file=sprintf("%s/iter/psi%d.RDS", base_path, s))
    saveRDS(0, file=sprintf("%s/iter/flag%d", base_path, s))
  } else {
    Sys.sleep(60)
  }
}


